"""
基于RSRS指标的多标的策略
使用分钟K线，多品种RSRS策略
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
from sklearn.linear_model import LinearRegression
from .base_strategy import BaseStrategy

class RSRStrategy(BaseStrategy):
    strategy_name = "rsrs"
    
    """
    基于RSRS指标的多标的策略
    使用分钟K线，多品种RSRS策略
    """
    def __init__(self, api: TqApi, symbol_info:dict, market_period:int):
        super().__init__(api, symbol_info, market_period)
        for symbol,info in symbol_info.items():
            self.all_info[symbol]["rsrs_slopes"] = []

        
    def calculate_signal(self, symbol: str,current_pos:int, kline: pd.DataFrame) -> float:
        """
        计算交易信号
        
        Args:
            symbol: 品种代码
            current_pos: 当前持仓
            kline: K线数据
            
        Returns:
            交易信号: 1-做多, -1-做空, 0-无信号
        """
        rsrs_score = self.calculate_rsrs(symbol, kline)
        
        long_threshold = self.all_info[symbol]["long_threshold"]
        short_threshold = self.all_info[symbol]["short_threshold"]
        long_close_threshold = self.all_info[symbol]["long_close_threshold"]
        short_close_threshold = self.all_info[symbol]["short_close_threshold"]
        if current_pos == 0:
            if rsrs_score > long_threshold:
                signal =  1  # 做多信号
            elif rsrs_score < short_threshold:
                signal = -1  # 做空信号
            else:
                signal = 0  # 无信号
        elif current_pos > 0:
            if rsrs_score < short_threshold:
                signal = -1
            elif rsrs_score < long_close_threshold:
                signal = 0  # 平多信号
            else:
                signal = 1  # 无信号
        elif current_pos < 0:
            if rsrs_score > long_threshold:
                signal = 1  # 做多信号
            elif rsrs_score > short_close_threshold:
                signal = 0  # 平空信号
            else:
                signal = -1
        return signal, rsrs_score

    
    def calculate_rsrs(self, symbol, kline) -> float:
        """
        计算RSRS指标
        
        Args:
            symbol: 品种代码
            kline: K线数据
            
        Returns:
            标准化后的RSRS值
        """
        rsrs_period = self.all_info[symbol]["rsrs_period"]
        rsrs_std_period = self.all_info[symbol]["rsrs_std_period"]
        
        # 获取最近的数据
        high_prices = kline.high.values
        low_prices = kline.low.values
        
        
        rsrs_slopes = self.all_info[symbol]["rsrs_slopes"]
        
        # 首次计算，从足够的历史数据开始计算
        if len(rsrs_slopes) < rsrs_std_period:
            for i in range(rsrs_period, len(high_prices)+1):
                # 获取计算窗口内的高低价
                window_high = high_prices[i-rsrs_period:i]
                window_low = low_prices[i-rsrs_period:i]

                # 线性回归计算斜率
                X = window_low.reshape(-1, 1)  # 最低价作为自变量
                y = window_high  # 最高价作为因变量
                
                reg = LinearRegression().fit(X, y)
                slope = reg.coef_[0]
                self.all_info[symbol]["rsrs_slopes"].append(slope)

        # 后续更新计算
        else:
            window_high = high_prices[-rsrs_period:]
            window_low = low_prices[-rsrs_period:]
            # 线性回归计算斜率
            X = window_low.reshape(-1, 1)  # 最低价作为自变量
            y = window_high  # 最高价作为因变量
            reg = LinearRegression().fit(X, y)
            slope = reg.coef_[0]
            self.all_info[symbol]["rsrs_slopes"][:-1] = self.all_info[symbol]["rsrs_slopes"][1:]
            self.all_info[symbol]["rsrs_slopes"][-1]=slope
        
        # 获取最新的RSRS斜率
        current_slope = self.all_info[symbol]["rsrs_slopes"][-1]
        
        # 标准化处理
        recent_slopes = self.all_info[symbol]["rsrs_slopes"][-rsrs_std_period:] 
        mean_slope = np.mean(recent_slopes)
        std_slope = np.std(recent_slopes)
        
        if std_slope > 0:
            rsrs_score = (current_slope - mean_slope) / std_slope
        else:
            rsrs_score = 0.0
        
        return rsrs_score
